import os
import music21 as m21
import json
import numpy as np
import tensorflow.keras as keras


KERN_DATASET_PATH = "D:\\Audio\\audio_datasets\\essen\\asia\\china\\shanxi"
ACCEPTABLE_DURATIONS = [
    0.25,
    0.5,
    0.75,
    1,
    1.5,
    2,
    3,
    4
]
SAVE_DIR = 'dataset_shanxi'
SINGLE_FILE_DATASET = 'file_dataset_shanxi'
SEQUENCE_LENGTH = 64
MAPPING_PATH = 'mapping_shanxi.json'


def load_songs_in_kern(dataset_path):
    
    songs = []
    
    # go through all the files in the dataset and load them with music21
    for path, subdirs, files in os.walk(dataset_path):
        for f in files:
            if f[-3:] == 'krn':
                # music21: kern, MIDI, MusicXML -> m21 -> kern, MIDI (converter)
                song = m21.converter.parse(os.path.join(path, f))
                songs.append(song)
    
    return songs


def has_acceptable_durations(song, acceptable_durations):
    """
    make DL more easier 去除奇异的拍数
    """
    for note in song.flat.notesAndRests:  # flatten to a single list.filter not notes and rests
        if note.duration.quarterLength not in acceptable_durations:
            return False
    return True
    

def transpose(song):
    """
    把所有歌统一转换成C大调
    """
    # get the key from the song(if has a key)
    # score: 整个片段  
    # parts: 声部  
    # measures: 小节，里面可能包含升降号（clef)，节奏型（第一小节）
    # notes: 音符
    parts = song.getElementsByClass(m21.stream.Part)
    measures_part0 = parts[0].getElementsByClass(m21.stream.Measure)
    key = measures_part0[0][4]
    
    # estimate key using music21?????
    if not isinstance(key, m21.key.Key):
        key = song.analyze('key')
    
    # print(key)
    
    # get interval for transposition. E.g., Bmaj -> Cmaj
    if key.mode == 'major':
        interval = m21.interval.Interval(key.tonic, m21.pitch.Pitch('C'))  # tonic: 
    elif key.mode == 'minor':
        interval = m21.interval.Interval(key.tonic, m21.pitch.Pitch('A'))
        
    # transpose song by calculated interval
    transposed_song = song.transpose(interval)
    
    return transposed_song
    

def encode_song(song, time_step=0.25):
    """
    把song编码成这种格式，一个step是1/4拍，_代表持续
    pitch = 60(middle C), duration = 1.0 -> [60, '_', '_', '_']
    """
    encoded_song = []
    
    for event in song.flat.notesAndRests:
        
        # handle notes 
        if isinstance(event, m21.note.Note):
            symbol = event.pitch.midi  # 60
        # handle rests 
        elif isinstance(event, m21.note.Rest):
            symbol = 'r' 
        
        # convert the note/rest into time series Notation
        steps = int(event.duration.quarterLength / time_step)
        for step in range(steps):
            if step == 0:
                encoded_song.append(symbol)
            else:
                encoded_song.append('_')
    # cast encoded song to a string
    encoded_song = ' '.join(map(str, encoded_song))
    
    return encoded_song

      
def preprocess(dataset_path):
    
    # load the folk songs
    print('------------------loading songs-------------------')
    songs = load_songs_in_kern(dataset_path)
    print(f"Loaded {len(songs)} songs.")
    
    for i, song in enumerate(songs):
        
        # filter out songs that have non-acceptable durations
        if not has_acceptable_durations(song, ACCEPTABLE_DURATIONS):
            continue
    
        # transpose songs to Cmaj/Amin
        song = transpose(song)
        
        # encode songs with music time series representation
        encoded_song = encode_song(song)
        
        # save songs to text files
        if not os.path.exists(SAVE_DIR):
            os.mkdir(SAVE_DIR)
        save_path = os.path.join(SAVE_DIR, str(i))
        
        with open(save_path, 'w') as fp:
            fp.write(encoded_song)
            fp.close()


def load(file_path):
    with open(file_path, "r") as fp:
        song = fp.read()
    
    return song


def create_single_file_dataset(dataset_path, file_dataset_path, sequence_length):
    """
    把dataset的文件归到一个文件里，方便LSTM训练
    """
    # 标记不同的song，长度为一个训练集的长度
    new_song_delimiter = '/ ' * sequence_length
    songs = ''
    
    # load encoded songs and add delimiters
    for path, _, files in os.walk(dataset_path):
        for f in files:
            file_path = os.path.join(path, f)
            song = load(file_path)
            songs = songs + song + ' ' + new_song_delimiter
    
    # remove the end space
    songs = songs[:-1]
    
    # save string that contains all dataset 
    with open(file_dataset_path, 'w') as fp:
        fp.write(songs)
        fp.close()
    
    return songs


def create_mapping(songs, mapping_path):
    """
    有些特殊符号，比如_或者/，LSTM只能接收数字，需要把他们map到整数上
    这里每一步都保存了文件，为了防止出错
    """
    mappings = {}
    
    # identify the vocabulary
    songs = songs.split()
    vocabulary = list(set(songs))
    
    # create mappings
    for i, symbol in enumerate(vocabulary):
        mappings[symbol] = i
    
    # save vocabulary to a json file
    with open(mapping_path, 'w') as fp:
        json.dump(mappings, fp, indent=4)
    

def convert_songs_to_int(songs):
    """
    用保存的mapping json文件，把每首歌的所有字符转换成数字
    """
    int_songs = []
    
    # load mappings
    with open(MAPPING_PATH, 'r') as fp:
        mappings = json.load(fp)
    
    # cast songs string to a list
    songs = songs.split()
    
    # map songs to int
    for symbol in songs:
        int_songs.append(mappings[symbol])
    
    return int_songs


def generate_sequences(sequence_length):
    """
    生成训练序列，用前面suquence_length长度的音符作为输入，后面一个音符作为输出
    最后用one-hot编码
    """
    # dataset: [11, 12, 13, 14, ...] -> inputs: [11, 12], target: 13; inputs: [12, 13], target: 14
    
    # load songs and map them to int
    songs = load(SINGLE_FILE_DATASET)
    int_songs = convert_songs_to_int(songs)
    
    # generate the training sequences
    # 100 symbols, 64 sequence length, 100 - 64 = 36 sequences
    inputs = []
    targets = []
    
    num_sequences = len(int_songs) - sequence_length
    for i in range(num_sequences):
        inputs.append(int_songs[i:i+sequence_length])
        targets.append(int_songs[i+sequence_length])
    
    # one-hot encode the sequences
    # inputs: (number of sequences, sequence length, vocabulary_size)
    # [ [0, 1, 2], [1, 1, 2] ] -> [ [ [1, 0, 0], [0, 1, 0], [0, 0, 1] ], [] ]
    vocabulary_size = len(set(int_songs))
    inputs = keras.utils.to_categorical(inputs, num_classes=vocabulary_size)
    targets = np.array(targets)
    
    return inputs, targets
    
    
if __name__ == "__main__":
    
    preprocess(KERN_DATASET_PATH)
    songs = create_single_file_dataset(SAVE_DIR, SINGLE_FILE_DATASET, SEQUENCE_LENGTH)
    create_mapping(songs, MAPPING_PATH)
    inputs, targets = generate_sequences(SEQUENCE_LENGTH)
